import os
import time
import base64
from abc import ABC, abstractmethod
from typing import Dict, Optional, List
from volcenginesdkarkruntime import Ark
from pydantic import BaseModel
from ..config.annotation import print_time
from loguru import logger

class VideoGenerationConfig(BaseModel):
    model: str
    api_key: Optional[str] = None
    base_url: Optional[str] = None


class VideoGenerationResult(BaseModel):
    task_id: str
    status: str = "queued"
    video_url: Optional[str] = None
    error: Optional[str] = None
    usage: Optional[Dict] = None
    created_at: float = None

    def __init__(self, **data):
        if 'created_at' not in data:
            data['created_at'] = time.time()
        super().__init__(**data)


class VideoGenerationListResult(BaseModel):
    items: List[VideoGenerationResult]
    total: int


class VideoGenerator(ABC):
    def __init__(self, config: VideoGenerationConfig):
        self.config = config

    @classmethod
    def build(cls, generator_type: Optional[str] = None, config: Optional[VideoGenerationConfig] = None) -> 'BaseVideoGenerator':
        """Factory method to create video generator based on type"""
        if generator_type is None:
            generator_type = os.getenv('LLM_VIDEO', 'volcengine')

        if config is None:
            config = cls._create_config_from_env(generator_type)

        generator_map = {'volcengine': 'VolcengineVideoGenerator'}

        generator_type_lower = generator_type.lower()
        if generator_type_lower not in generator_map:
            raise ValueError(f"Unsupported generator type: {generator_type}")

        class_name = generator_map[generator_type_lower]
        import sys
        current_module = sys.modules[__name__]

        if not hasattr(current_module, class_name):
            raise ValueError(f"Generator class {class_name} not found")

        generator_class = getattr(current_module, class_name)
        return generator_class(config)

    @classmethod
    def _create_config_from_env(cls, generator_type: str) -> VideoGenerationConfig:
        """Create configuration from environment variables"""
        prefix = generator_type.upper()

        model = os.getenv(f"{prefix}_VIDEO_MODEL")
        api_key = os.getenv(f"{prefix}_VIDEO_API_KEY")
        base_url = os.getenv(f"{prefix}_VIDEO_BASE_URL")

        if not model:
            raise ValueError(f"Environment variable {prefix}_VIDEO_MODEL is required")
        if not api_key:
            raise ValueError(f"Environment variable {prefix}_VIDEO_API_KEY is required")

        return VideoGenerationConfig(model=model, api_key=api_key, base_url=base_url)

    @abstractmethod
    def generate_text_to_video(self, text: str, callback_url: Optional[str] = None, **kwargs) -> VideoGenerationResult:
        pass

    @abstractmethod
    def generate_image_to_video(self, image: str, text: Optional[str] = None,
                                callback_url: Optional[str] = None, **kwargs) -> VideoGenerationResult:
        pass

    @abstractmethod
    def query(self, task_id: str) -> VideoGenerationResult:
        pass

    @abstractmethod
    def query_list(self, page_num: Optional[int] = None, page_size: Optional[int] = None,
                   status: Optional[str] = None, task_ids: Optional[List[str]] = None,
                   model: Optional[str] = None) -> VideoGenerationListResult:
        pass

    def _process_image(self, image: str) -> str:
        """Process image input and return data URL"""
        if image.startswith(('http://', 'https://', 'data:image/')):
            return image
        elif ',' in image and not os.path.exists(image):
            return f"data:image/jpeg;base64,{image}"
        else:
            with open(image, "rb") as f:
                image_data = base64.b64encode(f.read()).decode("utf-8")
            return f"data:image/jpeg;base64,{image_data}"

    def wait_for_completion(self,
                            task_id: str,
                            max_wait_time: int = 300,
                            poll_interval: int = 5) -> VideoGenerationResult:
        """Wait for video generation task to complete"""
        start_time = time.time()

        while time.time() - start_time < max_wait_time:
            result = self.query(task_id)

            if result.status in ["succeeded", "failed", "cancelled"]:
                return result
            elif result.status in ["queued", "running"]:
                time.sleep(poll_interval)
            else:
                return result

        return VideoGenerationResult(
            task_id=task_id,
            status="timeout",
            error=f"Task did not complete within {max_wait_time} seconds"
        )


class VolcengineVideoGenerator(VideoGenerator):
    """Volcano Engine video generation implementation"""

    def __init__(self, config: VideoGenerationConfig):
        super().__init__(config)
        self.client = Ark(api_key=config.api_key)

    def _build_parameters_text(self, **kwargs) -> str:
        """Build parameters text for Volcano video generation"""
        param_map = {
            'resolution': ('--rs', "720p"),
            'ratio': ('--rt', "adaptive"),
            'duration': ('--dur', 5),
            'fps': ('--fps', 24),
            'watermark': ('--wm', False),
            'seed': ('--seed', -1),
            'camera_fixed': ('--cf', False)
        }

        params = []
        for key, (flag, default) in param_map.items():
            value = kwargs.get(key, default)
            if key == 'seed' and value == -1:
                continue
            if key in ['watermark', 'camera_fixed']:
                value = str(value).lower()
            params.append(f"{flag} {value}")

        return " ".join(params)

    @print_time
    def generate_text_to_video(self, text: str,
                               callback_url: Optional[str] = None,
                               wait_for_completion: bool = True,
                               max_wait_time: int = 300,
                               **kwargs) -> VideoGenerationResult:
        """Generate video from text using Volcano Engine"""
        try:
            params_text = self._build_parameters_text(**kwargs)
            full_text = f"{text} {params_text}".strip()

            logger.info(f"Starting text-to-video generation with text: {full_text}")

            resp = self.client.content_generation.tasks.create(
                model=self.config.model,
                content=[{"text": full_text, "type": "text"}],
                callback_url=callback_url
            )

            logger.info(f"Video generation task created successfully, task_id: {resp.id}")
            result = VideoGenerationResult(task_id=resp.id, status="queued")

            if wait_for_completion:
                logger.info(f"Waiting for task completion, max_wait_time: {max_wait_time}s")
                return self.wait_for_completion(resp.id, max_wait_time)
            else:
                return result

        except Exception as e:
            logger.error(f"Failed to generate text-to-video: {str(e)}")
            return VideoGenerationResult(task_id="", status="failed", error=str(e))

    @print_time
    def generate_image_to_video(self, image: str, text: Optional[str] = None,
                                callback_url: Optional[str] = None,
                                wait_for_completion: bool = True,
                                max_wait_time: int = 300, **kwargs) -> VideoGenerationResult:
        """Generate video from image using Volcano Engine"""
        try:
            content = []

            if text:
                params_text = self._build_parameters_text(**kwargs)
                full_text = f"{text} {params_text}".strip()
                content.append({"text": full_text, "type": "text"})
                logger.debug(f"generate_image_to_video text content: {full_text}")

            image_url = self._process_image(image)
            content.append({"image_url": {"url": image_url}, "type": "image_url"})

            # Truncate image_url for logging if it's too long
            display_image_url = image_url[:100] + "..." if len(image_url) > 100 else image_url
            logger.info(f"Starting image-to-video generation with image: {display_image_url} , with text prompt: {full_text if full_text else t}")

            resp = self.client.content_generation.tasks.create(
                model=self.config.model,
                content=content,
                callback_url=callback_url
            )

            logger.info(f"Video generation task created successfully, task_id: {resp.id}")
            result = VideoGenerationResult(task_id=resp.id, status="queued")

            if wait_for_completion:
                logger.info(f"Waiting for task completion, max_wait_time: {max_wait_time}s")
                return self.wait_for_completion(resp.id, max_wait_time)
            else:
                return result

        except Exception as e:
            logger.error(f"Failed to generate image-to-video: {str(e)}")
            return VideoGenerationResult(task_id="", status="failed", error=str(e))

    def query(self, task_id: str) -> VideoGenerationResult:
        try:
            resp = self.client.content_generation.tasks.get(task_id=task_id)

            video_url = None
            error = None
            usage = None

            if resp.status == "succeeded" and hasattr(resp, 'content') and resp.content:
                video_url = getattr(resp.content, 'video_url', None)

            if resp.status == "failed" and hasattr(resp, 'error') and resp.error:
                error = getattr(resp.error, 'message', None) or f"Error code: {getattr(resp.error, 'code', 'unknown')}"

            if hasattr(resp, 'usage') and resp.usage and hasattr(resp.usage, 'completion_tokens'):
                usage = {"completion_tokens": resp.usage.completion_tokens}

            return VideoGenerationResult(
                task_id=task_id,
                status=resp.status,
                video_url=video_url,
                error=error,
                usage=usage
            )

        except Exception as e:
            return VideoGenerationResult(task_id=task_id, status="failed", error=str(e))

    def query_list(self, page_num: Optional[int] = None, page_size: Optional[int] = None,
                   status: Optional[str] = None, task_ids: Optional[List[str]] = None,
                   model: Optional[str] = None) -> VideoGenerationListResult:
        """Query video generation task list using Volcano Engine"""
        try:
            kwargs = {}
            if page_num is not None:
                kwargs['page_num'] = page_num
            if page_size is not None:
                kwargs['page_size'] = page_size
            if status is not None:
                kwargs['status'] = status
            if task_ids is not None:
                kwargs['task_ids'] = task_ids
            if model is not None:
                kwargs['model'] = model

            resp = self.client.content_generation.tasks.list(**kwargs)

            items = []
            for item in resp.items:
                video_url = None
                error = None
                usage = None

                if item.status == "succeeded" and hasattr(item, 'content') and item.content:
                    video_url = getattr(item.content, 'video_url', None)

                if item.status == "failed" and hasattr(item, 'error') and item.error:
                    error = getattr(item.error, 'message', None) or f"Error code: {getattr(item.error, 'code', 'unknown')}"

                if hasattr(item, 'usage') and item.usage and hasattr(item.usage, 'completion_tokens'):
                    usage = {"completion_tokens": item.usage.completion_tokens}

                items.append(VideoGenerationResult(
                    task_id=item.id,
                    status=item.status,
                    video_url=video_url,
                    error=error,
                    usage=usage,
                    created_at=getattr(item, 'created_at', time.time())
                ))

            return VideoGenerationListResult(items=items, total=resp.total)

        except Exception as e:
            return VideoGenerationListResult(items=[], total=0)
